package digitRecognitionProblem.learnWeights;

import genetic_algorithm.Chromosome;
import genetic_algorithm.Crossover;
import genetic_algorithm.FitnessFunction;
import genetic_algorithm.Mutation;
import genetic_algorithm.Population;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.io.PrintStream;
import java.util.LinkedList;
import java.util.List;

import mlp.Mlp;
import utils.GAretVal;
import digitRecognitionProblem.DigitRecognitionMlpFitness;

/**
 * Runs a genetic algorithm which finds the optimal initial weights of
 * a neural network which recognizes the specified digit.
 * The final network is serialized, and can be retrained using Retrain.java.
 * ATTENTION: before running, the activation function of mlp.Neuron.java must be set to ReLu
 * and mlp.NEGATIVE should be 0
 */
public class LWmain {

	// indices for main arguments
	private enum ArgNums {
		DIGIT, MAX_GENS, POP_SIZE, ELITE_SIZE, MUTATION_RATE, CROSSOVER_RATE, SINGLE_POINT_PROB, DOUBLE_POINT_PROB, UNIFORM_PROB, NUM_NEURONS, MLPS_DIR
	}
	
	private static int MAX_GENS = 50; // maximal number of generations to run	
	private static final int SELECTION_MODE = LWselection.ROULETTE; // selection phase mode
	private static final int TOURNAMENT_SIZE = 7; // how many chromosomes to include in tournament (if such mode is chosen for selection)
	private static int POP_SIZE = 25; // number of chromosomes in the population
	private static double CROSSOVER_RATE = 0.5; // chance to perform crossover
	private static double MUTATION_RATE = 0.1; // chance to perform mutation
	private static final String TEST_FILE = "validate1_reduce.csv"; // "test_reduce.txt"; // path to test file
	private static boolean REMOVE_PARENTS = false; // if chromosomes can be selected only once
	private static int ELITE_SIZE = 4; // number of elite chromosomes	
	public static int DIGIT = 4; // digit to learn
	private static String MLPS_DIRECTORY = "new_func_rand_28_4/" + DIGIT; // "/home/ofri/Desktop/MLPs/new_func_rand_28_4/" + DIGIT; // path to serialized neural networks
	
	public static void main(String[] args) {

		// redirect output to file
		DIGIT = Integer.parseInt(args[ArgNums.DIGIT.ordinal()]);
		try {
			File file = new File("LWmain_Details_rec_" + DIGIT + ".txt");
			FileOutputStream fos;
			fos = new FileOutputStream(file);
			PrintStream ps = new PrintStream(fos);
			System.setOut(ps);
			System.setErr(ps);
		} catch (FileNotFoundException e1) {
			e1.printStackTrace();
		}
		MLPS_DIRECTORY = args[ArgNums.MLPS_DIR.ordinal()] + "/" + DIGIT; // path to serialized neural networks of specified digit
		MAX_GENS = Integer.parseInt(args[ArgNums.MAX_GENS.ordinal()]);
		POP_SIZE = Integer.parseInt(args[ArgNums.POP_SIZE.ordinal()]);
		ELITE_SIZE = Integer.parseInt(args[ArgNums.ELITE_SIZE.ordinal()]);
		MUTATION_RATE = Double.parseDouble(args[ArgNums.MUTATION_RATE.ordinal()]);
		CROSSOVER_RATE = Double.parseDouble(args[ArgNums.CROSSOVER_RATE.ordinal()]);
		
		// document arguments
		System.out.println("digit = " + DIGIT);
		System.out.println("random mlps directory= " + MLPS_DIRECTORY);
		System.out.println("generations = " + MAX_GENS);
		System.out.println("population size = " + POP_SIZE);
		System.out.println("elite size = " + ELITE_SIZE);
		System.out.println("mutation rate = " + MUTATION_RATE);
		System.out.println("crossover rate = " + CROSSOVER_RATE);

		// initial weights range for each digit
		float[] initWeights = {0.0625f, 0.0625f, 0.078125f, 0.03125f, 0.09375f,  0.03125f, 0.125f, 0.09375f, 0.0625f, 0.140625f};
		
		// allocate objects to implement each phase of the genetic algorithm
		FitnessFunction fitnessFunction = new Lwfitness();
		DigitRecognitionMlpFitness.initTestData(TEST_FILE, DIGIT);		
		LW_Selection selector = new LWselection(REMOVE_PARENTS, SELECTION_MODE, TOURNAMENT_SIZE);		
		Crossover crossover = new LWcrossover();
		Mutation mutation = new LWmutation();
		
		// set parameters for the crossover
		LWcrossover.SINGLE_MODE_PROB = Float.parseFloat(args[ArgNums.SINGLE_POINT_PROB.ordinal()]);
		LWcrossover.DOUBLE_MODE_PROB = Float.parseFloat(args[ArgNums.DOUBLE_POINT_PROB.ordinal()]);
		LWcrossover.UNIFORM_MODE_PROB = Float.parseFloat(args[ArgNums.UNIFORM_PROB.ordinal()]);
		
		// set ranges for the mutation
		LWmutation.NUM_NEURONS = Integer.parseInt(args[ArgNums.NUM_NEURONS.ordinal()]);
		Mlp.INIT_WEIGHT_RANGE = initWeights[DIGIT];
		
		System.out.println("initializing population");
		
		// read serialized neural networks
		List<Mlp> mlpsList = readMlps();
		
		// create initial population
		List<Chromosome> chromList = new LinkedList<Chromosome>();
		for (int i = 0; i < POP_SIZE; ++i) {
			chromList.add(new LWchromosome(mlpsList.get(i)));
		}
		Population population = new Population(chromList, fitnessFunction);		
		
		System.out.println("population is set");
		
		// initialize the genetic algorithm
		LW_GeneticAlgorithm GA = new LW_GeneticAlgorithm(fitnessFunction, crossover,
				selector, mutation, CROSSOVER_RATE, MUTATION_RATE);
		
		// find optimal solution
		GAretVal result = GA.execute(population, MAX_GENS, true, ELITE_SIZE);
		
		// generate csv for execution data
		result.createsCsvFiles("LWmain_digit_" + DIGIT);
		
		FileOutputStream writeFile;
		try {
			writeFile = new FileOutputStream("mlp_0_rec_" + DIGIT + ".txt");
			ObjectOutputStream fileStream = new ObjectOutputStream(writeFile);			
			Mlp bestStart = ((LWchromosome) result.getOptSolution()).getMlp();
			fileStream.writeObject(bestStart);
			fileStream.close();
		} catch (FileNotFoundException e) {
			e.printStackTrace();
		} catch (IOException e) {
			e.printStackTrace();
		}
	}

	private static List<Mlp> readMlps() {

		List<Mlp> mlpsList = new LinkedList<Mlp>();
		File folder = new File(MLPS_DIRECTORY);
		for(File file : folder.listFiles())
		{
			System.out.print("|");
			try {
				FileInputStream readFile = new FileInputStream(file);
				ObjectInputStream fileStream = new ObjectInputStream(readFile);
				mlpsList.add((Mlp)fileStream.readObject());
				fileStream.close();
			} catch (Exception e) {
				e.printStackTrace();
			} 
		}
		System.out.println("reading mlps completed.");
		
		return mlpsList;
	}
}
